import sqlite3
import threading
from contextlib import contextmanager


from langgraph.checkpoint.sqlite import SqliteSaver


class ThreadSafeSqliteSaver(SqliteSaver):
    def __init__(self, db_path):
        super().__init__(db_path)
        self.db_path = db_path
        self.local = threading.local()

    def get_conn(self):
        if not hasattr(self.local, "conn"):
            self.local.conn = sqlite3.connect(self.db_path)
        return self.local.conn

    @contextmanager
    def cursor(self, transaction=True):
        conn = self.get_conn()
        cur = conn.cursor()
        try:
            yield cur
            if transaction:
                conn.commit()
        except Exception:
            if transaction:
                conn.rollback()
            raise
        finally:
            cur.close()